# docker run -itd --gpus all --log-driver none --restart unless-stopped --name yolo_api -v /home/bluelans/yolo_api:/usr/src/yolo_api -w /usr/src/yolo_api -p 60608:60608 ultralytics/ultralytics:latest python main.py
# docker run -itd --log-driver none --restart unless-stopped --name yolo_api -v /home/bluelans/yolo_api:/usr/src/yolo_api -w /usr/src/yolo_api -p 60608:60608 ultralytics/ultralytics:latest python main.py
import base64
from subprocess import check_output

import cv2, numpy
import requests
from ultralytics import YOLO

try:
    from fastapi import FastAPI
    from pydantic import BaseModel
except ModuleNotFoundError:
    check_output(["pip", "install", "fastapi", "-i", "https://mirrors.aliyun.com/pypi/simple"]).decode()
finally:
    from fastapi import FastAPI

try:
    import uvicorn
except ModuleNotFoundError:
    check_output(["pip", "install", "uvicorn", "-i", "https://mirrors.aliyun.com/pypi/simple"]).decode()
finally:
    import uvicorn

app = FastAPI()

ibytes_model = YOLO('ibytes.pt')
# ebay_model = YOLO('ebay.pt')
ebay_models = {
    "something you can write with": "ebay.pt",
    "a computer mouse": "ebay.pt",
    "a bumblebee": "a bumblebee.pt",
    "a plant": "a plant.pt",
}

sxf_model = YOLO("sxf_captcha.pt")
number_letters_model = YOLO("number_letters.pt")
yolov8s_model = YOLO("yolov8s.pt")
pingpong_model = YOLO("pingpong.pt")
infinia_model = YOLO("infinia_captcha.pt")


class IbytesCaptchaItem(BaseModel):
    base64_img: str
    conf: float = 0.5
    n: bool = True


class EbayCaptchaItem(BaseModel):
    base64_img: str
    conf: float = 0.5
    question: str


class SXFCaptcha(BaseModel):
    base64_img: str


def base64_to_im(b: str):
    img_bytes = base64.b64decode(b.encode())
    return bytes_to_im(img_bytes)


def bytes_to_im(b: bytes):
    return cv2.imdecode(numpy.frombuffer(b, numpy.uint8), 1)


class YoloV8sItem(BaseModel):
    base64_img: str = ""
    conf: float = 0.5
    label: str = "person"
    url: str = ""


class PingPongItem(BaseModel):
    base64_img: str = ""
    conf: float = 0.5


@app.post("/ibytes_captcha")
def ibytes_captcha(item: IbytesCaptchaItem):
    im = base64_to_im(item.base64_img)
    results = ibytes_model(im, conf=item.conf, device="cpu")
    for r in results:
        temp_list = {}
        indexes = []
        for i, v in enumerate(r.boxes.cls):
            if int(v) not in temp_list.keys():
                temp_list[int(v)] = i
            else:
                indexes = [temp_list[int(v)], i]
                break

        ret = []
        if indexes:
            for i in indexes:
                if item.n:
                    ret.append({
                        "x": float(r.boxes.xywhn[i][0]),
                        "y": float(r.boxes.xywhn[i][1]),
                        "w": float(r.boxes.xywhn[i][2]),
                        "h": float(r.boxes.xywhn[i][3]),
                    })
                else:
                    ret.append({
                        "x": float(r.boxes.xywh[i][0]),
                        "y": float(r.boxes.xywh[i][1]),
                        "w": float(r.boxes.xywh[i][2]),
                        "h": float(r.boxes.xywh[i][3]),
                    })
            return {"code": 0, "msg": "识别成功", "data": ret}
        return {"code": -1, "msg": "识别失败，请重试", "data": ret}
    return {"code": -1, "msg": "识别失败，请重试", "data": []}


@app.post("/ebay_captcha")
def ebay_captcha(item: EbayCaptchaItem):
    im = base64_to_im(item.base64_img)

    if not ebay_models.get(item.question):
        return {
            "code": -1,
            "message": f"{item.question}的模型未训练",
        }

    if not ebay_models.get(ebay_models[item.question]):
        ebay_models[ebay_models[item.question]] = YOLO(ebay_models[item.question])

    results = ebay_models[ebay_models[item.question]](im, conf=item.conf, device="cpu")
    for r in results:
        answer_index = [k for k, v in r.names.items() if v == item.question][0]
        results_index = [i for i, v in enumerate(r.boxes.cls) if int(v) == answer_index]

        ret = []
        if results_index:
            for r_ind in results_index:
                x = float(r.boxes.xywhn[r_ind][0]) // (1 / 3)
                y = float(r.boxes.xywhn[r_ind][1]) // (1 / 3)
                ret.append(int(x + 1 + y * 3))
            return {"code": 0, "msg": "识别成功", "data": list(set(ret))}
        return {"code": -1, "msg": "识别失败，请重试", "data": ret}
    return {"code": -1, "msg": "识别失败，请重试", "data": []}


@app.post("/sxf_captcha")
def sxf_captcha(item: SXFCaptcha):
    im = base64_to_im(item.base64_img)
    results = sxf_model(im)
    val = ""
    for r in results:
        temp_list = [(float(r.boxes.xyxy[i][0]), r.names[int(ni)]) for i, ni in enumerate(r.boxes.cls)]

        val = "".join([i[1] for i in sorted(temp_list, key=lambda k: k[0])])

    if len(val) == 4:
        return {"code": 0, "msg": "识别成功", "data": val}
    return {"code": -1, "msg": "识别失败，请重试", "data": None}


@app.post("/number_letters")
def sxf_captcha(item: SXFCaptcha):
    im = base64_to_im(item.base64_img)
    results = number_letters_model(im)
    val = ""
    for r in results:
        temp_list = [(float(r.boxes.xyxy[i][0]), r.names[int(ni)]) for i, ni in enumerate(r.boxes.cls)]
        val = "".join([i[1] for i in sorted(temp_list, key=lambda k: k[0])])
        return {"code": 0, "msg": "success", "data": val}
    return {"code": -1, "msg": "failed", "data": val}


@app.post("/yolov8s/predict")
def sxf_captcha(item: YoloV8sItem):
    if item.base64_img:
        im = base64_to_im(item.base64_img)
    elif item.url:
        resp = requests.get(item.url)
        content = resp.content
        im = bytes_to_im(content)
    else:
        return {"code": 1, "msg": f"没有图片信息", "data": None}

    results = yolov8s_model(im, conf=item.conf)
    results_index = []
    for r in results:
        answer_index = [k for k, v in r.names.items() if v == item.label][0]
        results_index = [i for i, v in enumerate(r.boxes.cls) if int(v) == answer_index]

    if len(results_index) > 0:
        return {"code": 0, "msg": f"图片包含{item.label}", "data": True}
    return {"code": 1, "msg": f"图片不包含{item.label}", "data": False}


@app.post("/pingpong_captcha")
def ibytes_captcha(item: PingPongItem):
    im = base64_to_im(item.base64_img)
    results = pingpong_model(im, device="cpu")

    for i, r in enumerate(results):
        if len(r.boxes.cls) == 2:
            pit_x = 0
            slice_x = 0
            for ind, i in enumerate(r.boxes.cls):
                if r.names[int(i)] == "slice":
                    slice_x = float(r.boxes.data[ind][0])
                elif r.names[int(i)] == "pit":
                    pit_x = float(r.boxes.data[ind][0])

            return {"code": 0, "msg": f"success", "data": pit_x - slice_x}
    return {"code": -1, "msg": f"failed", "data": 0}


@app.post("/infinia_captcha")
def sxf_captcha(item: SXFCaptcha):
    im = base64_to_im(item.base64_img)
    results = infinia_model(im)
    val = ""
    for r in results:
        temp_list = [(float(r.boxes.xyxy[i][0]), r.names[int(ni)]) for i, ni in enumerate(r.boxes.cls)]
        val = "".join([i[1] for i in sorted(temp_list, key=lambda k: k[0])])
        return {"code": 0, "msg": "success", "data": val}
    return {"code": -1, "msg": "failed", "data": val}

@app.get("/test")
async def test():
    return {'code': 0, "msg": "106图片识别服务"}


if __name__ == '__main__':
    uvicorn.run("main:app", port=60608, host='0.0.0.0', reload=False)
